from line_detect_img import *
import cv2
        
class VideoLineDetect():
    def __init__(self, video_src, output_save_file_path, if_paly_video = True, resize_width_height_tuple=()) -> None:
        self.video_src_path = video_src
        self.output_save_file_path = output_save_file_path
        self.if_paly_video = if_paly_video
        self.resize_width_height_tuple = resize_width_height_tuple
        self.vertices = None
        self.left_line_angle = -1
        self.right_line_angle = -1
        self.extrapolate_flag = True
        
    def set_line_RGB_color(self, line_RGB_color):
        # in the OpenCV, channel order is: BRG
        self.line_BGR_color = (line_RGB_color[2], line_RGB_color[1], line_RGB_color[0])
        
    def detect_line(self):
        if self.vertices is None:
            print("Vertices is None, please set the vertices!")
            return False
        
        if self.left_line_angle == -1 or self.right_line_angle == -1:
            print("angle is -1, please set the angle of the left and right lane line")
            return False
        
        cap = self.__get_capture__()
        if cap is None:
            print("get capture failed!")
            return False
        
        output_writer = self.__get_output_writer__(cap)
        if output_writer is None:
            print("Get the video output writer Failed!")
            return False
        
        while True:
            # 逐帧捕获
            ret, frame = cap.read()
            if not ret:
                # 到达视频末尾或出现错误
                print("End Of Stream Read!")
                break
            
            frame_raw, frame_resize = self.__frame_detec_line__(frame, 
                self.vertices, (self.left_line_angle, self.right_line_angle), self.extrapolate_flag)
                
            if output_writer is not None:
                # 将处理后的帧写入到输出视频
                output_writer.write(frame_raw)
            
            if self.if_paly_video:
                # 显示处理后的帧
                cv2.imshow('Processed Frame', frame_resize)
            
            # 按下q键退出
            if cv2.waitKey(1) & 0xFF == ord('q'):
                print("Capture the 'Q' key, so break loop!")
                break

        # 完成后释放资源
        cap.release()
        output_writer.release()
        cv2.destroyAllWindows()
        
        return True
        
    def set_verices(self, left_down_p, left_top_p, right_top, right_down):
        self.vertices = (left_down_p, left_top_p, right_top, right_down)
        
    def set_line_angle(self, pLeft_line_angle, pRight_line_angle):
        '''
        设置左右车道线角度，目的是构建左右车道线构成的三角形角度最大顶角，任何被识别出来的车道线只要没有在
        这个三角形范围内即可认为是有效识别；如果车道线在该三角形范围内，比如左车道线是30°，如果有某条被识别
        线段角度小于30°，那么就认为是在构建的三角形范围内，于是是无效线段
        
        比较好理解的说法是：指定左车道线的最小角度，右车道线的最大角度
        '''
        self.left_line_angle = pLeft_line_angle
        self.right_line_angle = pRight_line_angle
        
    def set_extrapolate_flag(self, extrapolate_flag):
        self.extrapolate_flag = extrapolate_flag
        
    def __get_capture__(self):
        # 打开视频文件或摄像头
        cap = cv2.VideoCapture(self.video_src_path)

        # 检查是否成功打开
        if not cap.isOpened():
            print("open video src ERROR! return")
            cap = None
            
        return cap
    
    def __get_output_writer__(self, cap):
        # 获取视频的属性
        fps = int(cap.get(cv2.CAP_PROP_FPS))
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

        # 定义输出的视频文件
        fourcc = cv2.VideoWriter_fourcc(*'MP4V')  # 可根据需要选择合适的编码器
        if self.output_save_file_path is not None and len(self.output_save_file_path) > 0:
            writer = cv2.VideoWriter(self.output_save_file_path, fourcc, fps, (width, height))
        else:
            writer = None
            
        return writer
    
    def __frame_detec_line__(self, frame, vertices, line_angle_tuple, extrapolate_flag):
        ld = Line_Detech()
        ld.set_verices(vertices[0], vertices[1], vertices[2], vertices[3])
        ld.set_line_angle(line_angle_tuple[0], line_angle_tuple[1])
        ld.set_show_image_flag(False)
        ld.set_line_color(self.line_BGR_color)
            
        processed_frame_raw = ld.detect_line_by_imgobj(frame, extrapolate_flag)
        if len(self.resize_width_height_tuple) > 0:
            processed_frame_resize = cv2.resize(processed_frame_raw, self.resize_width_height_tuple)
        else:
            processed_frame_resize = processed_frame_raw
            
        return processed_frame_raw,processed_frame_resize

        
    def sample(self, input_video_path):
        '''
        获取采样的一帧图像，用于分析IoR以及范围等
        '''
        cap = cv2.VideoCapture(input_video_path)
        
        # 检查是否成功打开
        if not cap.isOpened():
            print("无法打开视频源")
            exit()
            
        ret, frame = cap.read()
        if ret:
            save_path = "images/video_sample.jpg"
            cv2.imwrite(save_path, frame)

def main():
    input_video_path = "video/line_detect_5s.mp4"
    # 依次指定：
    # 1. 图像路径
    # 2. 保存路径
    # 3. 是否实时播放识别图像
    # 4. 如果指定实时播放，指定播放图像大小
    vld = VideoLineDetect(input_video_path, "video/test.mp4", True, (1980, 1080))
    # 设置感兴趣区域IoR（Interesting Of Region），顺序依次是：左下，左上，右上，右下
    vld.set_verices((0, 1750), (500, 1500), (3500, 1500), (4000, 1750))
    # 设置左右车道线角度范围
    vld.set_line_angle(20, 160)
    # 是否外推识别的车道线为直线（不构成直线的车道线段，将会被抛弃显示）
    vld.set_extrapolate_flag(True)
    # 设置描绘车道线段的颜色
    vld.set_line_RGB_color([255, 0, 0])
    
    # 识别车道线
    vld.detect_line()
    
if __name__ == "__main__":
    main()
    # sample()